from llama_index.embeddings.huggingface import HuggingFaceEmbedding
from llama_index.core import VectorStoreIndex, SimpleDirectoryReader
from llama_index.core.llms import CustomLLM, CompletionResponse, CompletionResponseGen
from llama_index.core.llms import LLMMetadata
from transformers import pipeline, AutoModelForCausalLM, AutoTokenizer
from typing import Any
from pydantic import Field

"""该代码实现了一个基于 RAG（检索增强生成） 的本地问答系统，专门针对 山西文旅 领域文档进行问答。核心功能包括：
文档检索：从本地文档中查找与问题相关的信息
答案生成：使用本地部署的 Qwen 大模型生成自然语言回答
完全本地化运行：不依赖 OpenAI 等在线 API（嵌入模型和 LLM 均在本地运行）
步骤	    代码	                                        说明
加载文档	SimpleDirectoryReader("山西文旅").load_data()	读取指定目录下的所有文档（支持 PDF/TXT 等格式）
文本嵌入	HuggingFaceEmbedding("bge-small-zh-v1.5")	使用本地 BGE 中文模型将文本转换为向量
构建索引	VectorStoreIndex.from_documents()	        创建可快速检索的向量数据库
"""

# 1. 嵌入模型
embed_model = HuggingFaceEmbedding("D:/ideaSpace/MyPython/models/bge-small-zh-v1.5")

# 2. 文档索引
documents = SimpleDirectoryReader("D:\\ideaSpace\\rag-in-action-master\\90-文档-Data\\山西文旅").load_data()
index = VectorStoreIndex.from_documents(documents, embed_model=embed_model)

# 3. 自定义本地 LLM
class LocalQwenLLM(CustomLLM):
    tokenizer: Any = Field(default=None, exclude=True)  # 声明为Pydantic字段
    model: Any = Field(default=None, exclude=True)      # 声明为Pydantic字段
    pipe: Any = Field(default=None, exclude=True)       # 声明为Pydantic字段

    def __init__(self, **kwargs):
        super().__init__(**kwargs)  # 必须调用父类初始化

        # 加载模型和tokenizer
        model_path = r"D:\ideaSpace\MyPython\models\qwen\Qwen1.5-0.5B-Chat"
        self.tokenizer = AutoTokenizer.from_pretrained(
            model_path,
            trust_remote_code=True
        )
        self.model = AutoModelForCausalLM.from_pretrained(
            model_path,
            device_map="auto",
            trust_remote_code=True
        )

        # 创建pipeline
        self.pipe = pipeline(
            "text-generation",
            model=self.model,
            tokenizer=self.tokenizer,
            max_new_tokens=256,
            temperature=0.1,
        )

    @property
    def metadata(self) -> LLMMetadata:
        """返回LLM元数据"""
        return LLMMetadata(
            model_name="Qwen1.5-0.5B-Chat",
            context_window=2048,
            num_output=256,
            is_chat_model=True
        )

    def complete(self, prompt: str, **kwargs: Any) -> CompletionResponse:
        """同步生成文本"""
        response = self.pipe(prompt)[0]["generated_text"]
        return CompletionResponse(text=response)

    def stream_complete(self, prompt: str, **kwargs: Any) -> CompletionResponseGen:
        """流式生成文本（简单实现）"""
        response = self.pipe(prompt)[0]["generated_text"]
        yield CompletionResponse(text=response, delta=response)

# 初始化LLM
llm = LocalQwenLLM()

# 4. 查询
query_engine = index.as_query_engine(llm=llm)
print(query_engine.query("山西省的主要旅游景点有哪些？"))